/*
 * (c) Copyright 2020 CORSIKA Project, corsika-project@lists.kit.edu
 *
 * This software is distributed under the terms of the 3-clause BSD license.
 * See file LICENSE for a full version of the license.
 */

/*
 * Squares3.hpp
 *
 *  Created on: 25/02/2021
 *      Author: Antonio Augusto Alves Junior
 */

#pragma once

#include <stdint.h>
#include "SquaresKeys.hpp"

namespace random_iterator {

  namespace detail {

    /*
     *  Three round counter-based middle square
     *
     *  squares3 - returns a 32-bit unsigned int [0,0xffffffff]
     *
     *
     *  Three rounds of squaring are performed and the result is returned.
     *  For the first two rounds, the result is rotated right 32 bits.
     *  This places the random data in the best position for the next round.
     *  y = ctr*key or z = (ctr+1)*key is added on each round.  For keys
     *  generated by the key utility, either ctr*key or (ctr+1)*key will
     *  have non-zero digits.  This improves randomization and also provides
     *  for a uniform output.
     *
     *  Note:  The squares RNG was based on ideas derived from Middle Square
     *  Weyl Sequence RNG.  One of the ideas was to obtain uniformity by adding
     *  the Weyl sequence after squaring.  Richard P. Brent (creator of the
     *  xorgens RNG) suggested this method.  It turns out that adding ctr*key
     *  is equivalent.  Brent's idea provides the basis for uniformity in this
     *  generator.
     *
     *  Implementation of the algorithm authored by Bernard Widynski and
     *  described in https://arxiv.org/pdf/2004.06278v2.pdf
     */

    class Squares3_64 {

    public:
      typedef uint64_t init_type;
      typedef uint64_t state_type;
      typedef uint64_t seed_type;
      typedef uint64_t advance_type;
      typedef uint32_t result_type;

      Squares3_64() = delete;

      Squares3_64(seed_type s, uint32_t)
          : state_(0)
          , seed_(seed_type{splitmix<seed_type>(s)}) {}

      Squares3_64(Squares3_64 const& other)
          : state_(other.getState())
          , seed_(other.getSeed()) {}

      inline Squares3_64& operator=(Squares3_64 const& other) {
        if (this == &other) return *this;

        state_ = other.getState();
        seed_ = other.getSeed();

        return *this;
      }

      inline result_type operator()(void) {
        uint64_t x, y, z;

        y = x = seed_ * state_;
        z = y + seed_;

        x = x * x + y;
        x = (x >> 32) | (x << 32); /* round 1 */

        x = x * x + z;
        x = (x >> 32) | (x << 32); /* round 2 */

        ++state_; /* advance state */

        return (x * x + y) >> 32; /* round 3 */
      }

      inline void discard(advance_type n) { state_ += n; }

      inline seed_type getSeed() const { return seed_; }

      inline void setSeed(seed_type seed) { seed_ = seed; }

      inline state_type getState() const { return state_; }

      inline void setState(state_type state) { state_ = state; }

      inline static uint64_t generateSeed(size_t i) { return keys[i]; }

      static constexpr result_type min() { return 0; }

      static constexpr result_type max() {
        return std::numeric_limits<result_type>::max();
      }

      friend inline std::ostream& operator<<(std::ostream& os, const Squares3_64& be) {
        return os << "state: " << be.getState() << " seed: " << be.getSeed();
      }

    private:
      state_type state_;
      seed_type seed_;
    };

  } // namespace detail

} // namespace random_iterator
